import xgboost as xgb
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
import matplotlib.pyplot as plt

"""
XGBoost 多分类示例 - 使用鸢尾花数据集
鸢尾花数据集包含3个类别：0-setosa, 1-versicolor, 2-virginica
特征数：4个（花萼长度、花萼宽度、花瓣长度、花瓣宽度）
样本数：150个
"""

# 加载数据并划分
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
X_train_final, X_val, y_train_final, y_val = train_test_split(
    X_train, y_train, test_size=0.2, random_state=42, stratify=y_train
)

print(f"数据形状: X{X.shape}, y{y.shape}")
print(f"类别分布: {np.bincount(y)}")


def train_model(
    X_train,
    y_train,
    X_val,
    y_val,
    max_depth=3,
    n_estimators=100,
    early_stopping_rounds=10,
):
    """训练XGBoost多分类模型"""
    model = xgb.XGBClassifier(
        objective="multi:softprob",
        max_depth=max_depth,
        n_estimators=n_estimators,
        random_state=42,
        eval_metric=["mlogloss"],
        early_stopping_rounds=early_stopping_rounds,
    )

    # 训练模型
    if X_val is not None and y_val is not None:
        # 使用验证集进行早停
        eval_set = [(X_train, y_train), (X_val, y_val)]
        model.fit(X_train, y_train, eval_set=eval_set, verbose=True)
    else:
        # 不使用验证集
        model.fit(X_train, y_train, verbose=True)

    # 获取训练历史
    if hasattr(model, "evals_result_"):
        print(f"最佳迭代次数: {model.best_iteration}")
        print(f"最佳验证损失: {model.best_score:.4f}")

    return model


def plot_history(model):
    """绘制训练历史"""
    results = model.evals_result_
    epochs = range(len(results["validation_0"]["mlogloss"]))

    fig, ax1 = plt.subplots(1, 1, figsize=(6, 4))

    # 损失曲线
    ax1.plot(epochs, results["validation_0"]["mlogloss"], label="Train")
    ax1.plot(epochs, results["validation_1"]["mlogloss"], label="Validation")
    ax1.set_title("Log Loss")
    ax1.legend()
    ax1.grid(True)

    plt.tight_layout()
    plt.show()


def evaluate_model(model, X_test, y_test):
    """评估模型性能"""
    # 获取概率预测
    y_pred_proba = model.predict_proba(X_test)
    # 获取类别预测
    y_pred = model.predict(X_test)

    accuracy = accuracy_score(y_test, y_pred)
    print(f"测试准确率: {accuracy:.4f}")

    print("前5个样本预测:")
    for i in range(min(5, len(y_pred_proba))):
        print(f"样本{i}: 真实={y_test[i]}, 预测={y_pred[i]}, 概率={y_pred_proba[i]}")

    return accuracy


def save_model(model, path="iris_xgboost_classifier.json"):
    """保存模型"""
    if not path.endswith(".json"):
        path += ".json"
    model.save_model(path)
    print(f"模型已保存: {path}")


if __name__ == "__main__":
    model = train_model(
        X_train_final,
        y_train_final,
        X_val,
        y_val,
    )
    plot_history(model)
    print("模型评估结果:")
    evaluate_model(model, X_test, y_test)
    # save_model(model)
